package com.taint.agent;


import com.taint.advice.SinkMethodAdvice;
import com.taint.advice.SourceMethodAdvice;
import com.taint.asm.core.ClassAnaylzer;
import com.taint.sourcesink.MethodHook;
import com.taint.utils.Utils;
import net.bytebuddy.agent.builder.AgentBuilder;
import net.bytebuddy.agent.builder.ResettableClassFileTransformer;
import net.bytebuddy.asm.Advice;
import net.bytebuddy.asm.AsmVisitorWrapper;
import net.bytebuddy.description.method.MethodDescription;
import net.bytebuddy.description.type.TypeDescription;
import net.bytebuddy.dynamic.DynamicType;
import net.bytebuddy.matcher.ElementMatcher;
import net.bytebuddy.matcher.ElementMatchers;
import net.bytebuddy.utility.JavaModule;
import org.apache.log4j.Logger;

import java.io.IOException;
import java.lang.instrument.Instrumentation;
import java.lang.instrument.UnmodifiableClassException;
import java.util.*;

import static com.taint.utils.HookType.*;
import static net.bytebuddy.matcher.ElementMatchers.*;

public class FlowTaintClassTransformer {
    static Logger logger = Logger.getLogger(FlowTaintClassTransformer.class);
    private static final boolean DEBUG = true;
    private  static HashMap<String, Boolean> isMatchMethod = new HashMap<>();
    private static Advice advice;
    private static HashMap<String, Boolean> isMatchParameter = new HashMap<>();

    public static ResettableClassFileTransformer getTransformer() {
        return transformer;
    }

    private static ResettableClassFileTransformer transformer;

    public static void ForwardStaticAnaylzer(Instrumentation instrumentation, HashMap<String, Boolean> matchedLoadedClazz){

        transformer =getInstance().type(nameStartsWith("com.huawei").or(nameStartsWith("org.joychou"))).transform(new AgentBuilder.Transformer() {
            @Override
            public DynamicType.Builder<?> transform(DynamicType.Builder<?> builder, TypeDescription typeDescription, ClassLoader classLoader, JavaModule module) {
                String clazzName=typeDescription.getCanonicalName();
                if(clazzName!=null && !clazzName.contains("$")){
                    matchedLoadedClazz.put(clazzName, true);
                    if (DEBUG) {
                        byte[] bytes = builder.make().getBytes();

                        if (bytes.length > 0) {
                            try {
                                Utils.OutPutClazz(bytes, typeDescription.getCanonicalName());
                                ClassAnaylzer.EntryClassAnaylzer(bytes, typeDescription.getCanonicalName());
                                  } catch (IOException e) {
                                e.printStackTrace();
                            }
                        }
                    }
                }
                return builder;
            }
        }).installOn(instrumentation);

        retransform(instrumentation);
    }

    public static AgentBuilder getInstance() {
        return new AgentBuilder.Default().ignore(
                        nameStartsWith("net.bytebuddy.")
                                .or(nameStartsWith("org.slf4j."))
                                .or(nameStartsWith("sun.reflect.")))
                // 声明使用redefine
                .with(AgentBuilder.RedefinitionStrategy.REDEFINITION)
                .with(AgentBuilder.TypeStrategy.Default.REDEFINE);
        //.with(AgentBuilder.Listener.StreamWriting.toSystemOut().withTransformationsOnly());
    }


    public static void InstrumenterSourceSink(String hooktype,
                                              LinkedHashMap<String, Set<MethodHook>> methodHookMap, Instrumentation inst) {

        String clazzName = null;
        if (hooktype.equals(SOURCE)) {
            advice = Advice.to(SourceMethodAdvice.class);
        } else if (hooktype.equals(SINK)) {
            advice = Advice.to(SinkMethodAdvice.class);
        }
        //循环
        //遍历效率高https://www.cnblogs.com/stono/p/4726177.html
        for (Map.Entry<String, Set<MethodHook>> stringSetEntry : methodHookMap.entrySet()) {
            //循环set
            for (Object o : (Set) ((Map.Entry) stringSetEntry).getValue()) {
                MethodHook methodHook = (MethodHook) o;
                String clazzType = methodHook.getType() != null ? methodHook.getType() : "default";
                clazzName = methodHook.getClazzName();
                String methodName = methodHook.getMethodName();
                String returnType = methodHook.getReturntype() != null ? methodHook.getReturntype() : "void";
                ArrayList<String> parameters = methodHook.getParameters() != null ? methodHook.getParameters() : null;
                if (null != methodName && methodName.equalsIgnoreCase(MATCHALL)) {
                    isMatchMethod.put(clazzName, true);
                }
                if (null != parameters && parameters.get(0).equalsIgnoreCase(MATCHALL)) {
                    isMatchParameter.put(clazzName, true);
                }

                transformer = getInstance().type(classMatch(clazzName, clazzType)).transform((builder, typeDescription, classLoader, module) -> {
                    // System.out.println("test: " + typeDescription.getCanonicalName());
                    String descriptor;
                    //method 开启全匹配
                    if (isMatchMethod.containsKey(typeDescription.getCanonicalName())) {
                        for (int k = 0; k < typeDescription.getDeclaredMethods().size(); k++) {
                            builder = builder.visit(new AsmVisitorWrapper.ForDeclaredMethods().method(named(typeDescription.getDeclaredMethods().get(k).getName()), Advice.to(SinkMethodAdvice.class)));
                        }
                    }
                    //method的参数 全匹配
                    else if (isMatchParameter.containsKey(typeDescription.getCanonicalName())) {
                        //System.out.println("method: " + methodName);
                        builder = builder.visit(new AsmVisitorWrapper.ForDeclaredMethods().method(methodMatch(methodName, MATCHALL), advice));
                    }
                    //精准匹配
                    else {
                        if (null == parameters) {
                            descriptor = Utils.toInternalMethodSignature(returnType, null);
                        } else {
                            descriptor = Utils.toInternalMethodSignature(returnType, parameters);
                        }
                        builder = builder.visit(new AsmVisitorWrapper.ForDeclaredMethods().method(methodMatch(methodName, descriptor), Advice.to(SinkMethodAdvice.class)));
                    }
                    if (DEBUG) {
                        byte[] bytes = builder.make().getBytes();

                        if (bytes.length > 0) {
                            try {
                               // ASM.ClazzBytes(bytes);
                                Utils.OutPutClazz(bytes, typeDescription.getCanonicalName());
                            } catch (IOException e) {
                                e.printStackTrace();
                            }
                        }
                    }
                    return builder;
                }).installOn(inst);
            }
            retransform(inst, clazzName);
        }

       // retransform(inst);
    }


    //https://github.com/bsyonline/AlphaX/blob/60a80404225d4a4f6f7b1f663462712c5b0d8820/byte-code/byte-buddy/auth/auth-agent/src/main/java/com/rolex/alphax/auth/agent/AuthAgent.java
    public static ElementMatcher<TypeDescription> classMatch(String clazzName, String clazzType) {
        ElementMatcher<TypeDescription> clazzMatch = null;
        if (clazzType.equalsIgnoreCase(INTERFACE)) {
            clazzMatch = not(isAbstract()).and(ElementMatchers.hasSuperType(named(clazzName)));
        } else if (clazzType.equalsIgnoreCase(ANNOTATION)) {
            clazzMatch = ElementMatchers.isAnnotatedWith(ElementMatchers.anyOf(named(clazzName)));
        } else if (clazzType.equalsIgnoreCase(ENUM)) {
            clazzMatch = ElementMatchers.isEnum().and(named(clazzName));
        } else {
            clazzMatch = not(isAbstract()).and(any()).and(named(clazzName));
        }

        return clazzMatch;
    }

    public static ElementMatcher<MethodDescription> methodMatch(String matchedMethod, String descriptor) {

        if (null == matchedMethod || null == descriptor) {
            return not(isSetter()).or(not(isGetter()));
        }
        if (descriptor.equalsIgnoreCase(MATCHALL)) {
            //nameMatches() 后续可以的支持正则
            return not(isSetter()).or(not(isGetter())).and(named(matchedMethod));
        } else {
            return not(isSetter()).or(not(isGetter())).and(named(matchedMethod).and(hasDescriptor(descriptor)));
        }

    }

    public static void retransform(Instrumentation inst) {
        //这个效率应该比较低，这里后面需要修改成Map

        Class[] loadedClasses = inst.getAllLoadedClasses();

        for (Class clazz : loadedClasses) {
            //logger.info("loaded class: "+clazz.getName());
            if (inst.isModifiableClass(clazz) && clazz.getName().startsWith("org.joychou")) {
                logger.info("all matched class: " + clazz.getName());
                try {
                    inst.retransformClasses(clazz);
                } catch (UnmodifiableClassException e) {
                    e.printStackTrace();
                }
            }

        }

    }

    public static void retransform(Instrumentation inst, String clazzName) {
        //这个效率应该比较低，这里后面需要修改成Map

        Class[] loadedClasses = inst.getAllLoadedClasses();

        for (Class clazz : loadedClasses) {
            //logger.info("loaded class: "+clazz.getName());
            if (inst.isModifiableClass(clazz) && clazz.getName().equals(clazzName)) {
                logger.info("matched class: " + clazz.getName());
                try {
                    inst.retransformClasses(clazz);
                } catch (UnmodifiableClassException e) {
                    e.printStackTrace();
                }
            }

        }

    }

}
